import json
import faiss
import numpy as np
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 将一维的文本数据通过模型转成二维向量
model_id = "nlp_corom_sentence-embedding_chinese-base-medical"
pipeline_se = pipeline(Tasks.sentence_embedding, model=model_id)

def retrievalVector(index: str):
    loaded_index = faiss.read_index(index)
    inputs = {
        'source_sentence': ["n香港大学深圳医院"],
    }
    result = pipeline_se(input=inputs)
    embeddings = result.get("text_embedding")
    print(embeddings)

    embeddings_array = np.array(embeddings).astype("float32").reshape(1, -1)  # 变成形状 (1, d)

    # 正确调用 search 方法
    I, D = loaded_index.search(embeddings_array, 5)  # 第一个参数是查询向量
    print("-" * 100)
    print(I)
    print("-" * 100)
    print(D)
    print("-" * 100)

    print(I[0].tolist())